import os
import json
import sys
import argparse
import xlwt


def config():
	parser = argparse.ArgumentParser()
	parser.add_argument("--query_list", type=str, required=True)
	parser.add_argument("--stra_input", type=str, required=True)
	parser.add_argument("--base_input", type=str, default="outputs")
	parser.add_argument("--output_file", type=str, default="bad_case_ana.xls")
	return parser.parse_args()

args = config()

# query list
with open(args.query_list) as f:
	query_list = f.read().strip().split("\n")

# stra input
stra_input = {}
with open(args.stra_input) as f:
	lines = f.read().strip().split("\n")
	for line in lines:
		obj = json.loads(line)
		stra_input[obj["query"]["Q_query_basic"]] = obj


def load_base(filename):
	with open(filename) as f:
		content = f.read().strip().split("\n")

	table = []
	for line in content:
		table.append(line.split("\t"))

	header = table[0]
	query_type_col = header.index("query_type")
	url_col = header.index("url")
	title_col = header.index("title")

	video_query_type = "QuerySemanticDwellingTime"
	text_query_type = "DwellTimeTextVecRecallQueue"

	max_u_num = 50
	v_content = []
	t_content = []

	for row in table[1:]:
		if (row[query_type_col] == video_query_type) and (len(v_content) < max_u_num):
			v_content.append([row[url_col], row[title_col]])
		elif (row[query_type_col] == text_query_type) and (len(t_content) < max_u_num):
			t_content.append([row[url_col], row[title_col]])

	return v_content, t_content


def load_stra(query):
	stra_item = stra_input[query]
	candidates_top = stra_item["candidates_top"]
	stra_cnt = []
	for item in candidates_top:
		stra_cnt.append("{}({:.6f})".format(item["doc"]["PU_title_basic"].replace("_", ""), item["sim"]))
	return stra_cnt


def align_and_cat(b_v_cnt, b_t_cnt, stra_cnt):
	mxl = 0
	if len(b_v_cnt) > mxl:
		mxl = len(b_v_cnt)
	if len(b_t_cnt) > mxl:
		mxl = len(b_t_cnt)
	if len(stra_cnt) > mxl:
		mxl = len(stra_cnt)

	b_v_cnt = b_v_cnt + [["-", "-"] for _ in range(mxl - len(b_v_cnt))]
	b_t_cnt = b_t_cnt + [["-", "-"] for _ in range(mxl - len(b_t_cnt))]
	stra_cnt = stra_cnt + ["-" for _ in range(mxl - len(b_t_cnt))]

	return [["基线-视频URL", "基线-视频TITLE", "基线-图文URL", "基线-图文TITLE", "策略"]] + \
		[i1 + i2 + [i3] for i1, i2, i3 in zip(b_v_cnt, b_t_cnt, stra_cnt)]


workbook = xlwt.Workbook()
style = xlwt.easyxf('align: wrap yes')
header_style = xlwt.easyxf('font: bold on; pattern: pattern solid, fore_colour gray25;')  # 创建一个首行的样式对象

for qidx, query in enumerate(query_list):
	print("process {}.{}".format(qidx, query))

	try:
		b_v_cnt, b_t_cnt = load_base(os.path.join(args.base_input, "{}.txt".format(query)))
		print(len(b_v_cnt), len(b_t_cnt))

		stra_cnt = load_stra(query)
	except Exception as e:
		continue

	worksheet = workbook.add_sheet("{}.{}".format(qidx, query))
	cat_cnt = align_and_cat(b_v_cnt, b_t_cnt, stra_cnt)

	for i, row in enumerate(cat_cnt):
		for j, value in enumerate(row):
			if i == 0:
				worksheet.write(i, j, value, header_style)
			else:
				worksheet.write(i, j, value, style)

	worksheet.col(0).width = 10000
	worksheet.col(1).width = 10000
	worksheet.col(2).width = 10000
	worksheet.col(3).width = 10000
	worksheet.col(4).width = 10000

workbook.save(args.output_file)
